{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import zipfile\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import akshare as ak\n",
    "\n",
    "## utility\n",
    "\n",
    "# 获取沪深300成分股\n",
    "hs300_stock_list = ak.index_stock_cons(symbol=\"000300\")\n",
    "stock_list = hs300_stock_list['品种代码']\n",
    "hs300_list = list(stock_list)\n",
    "\n",
    "## 获取所有股票上市时间\n",
    "def get_sh_time():    \n",
    "    df_sh = pd.read_excel('../量化data/股票代码.xlsx')\n",
    "    df_sh = df_sh.iloc[0:5355,:]\n",
    "    df_sh = df_sh.rename(columns={'证券代码↑':'code'})\n",
    "    df_sh['上市日期'] = pd.to_datetime(df_sh['上市日期']) \n",
    "    return df_sh\n",
    "\n",
    "# 1 MAD:中位数去极值\n",
    "def extreme_MAD(dt,n = 3):\n",
    "    \"\"\"\n",
    "    args:\n",
    "        dt:一列或多列时间序列,df\n",
    "        n : 中位数+-n*偏差中位数\n",
    "    return:\n",
    "        返回替换极值dt\n",
    "    \"\"\"\n",
    "    median = dt.quantile(0.5)   # 找出中位数\n",
    "    new_median = (abs((dt - median)).quantile(0.5))   # 偏差值的中位数\n",
    "    dt_up = median + n*new_median    # 上限\n",
    "    dt_down = median - n*new_median  # 下限\n",
    "    return dt.clip(dt_down, dt_up, axis=1)    # 超出上下限的值，赋值为上下限\n",
    "\n",
    "# 2 填补缺失值\n",
    "def Na_Fill(df,method='ffill'):\n",
    "    \"\"\"\n",
    "    args:\n",
    "        df:df格式数据\n",
    "        method: 填补方式，'ffill'前一个非NA值填充,'bfill'后一个非NA填充\n",
    "    return:\n",
    "        df\n",
    "    \"\"\"\n",
    "    df = df.fillna(method)\n",
    "    return df\n",
    "\n",
    "def get_sh_process(x,df_sh):\n",
    "    \"\"\"\n",
    "    去除上市前的样本\n",
    "    args:\n",
    "        x: 单只股票的df\n",
    "        df_sh:记录股票上市时间的df\n",
    "    return:\n",
    "        单只股票去除上市前样本的df\n",
    "    \"\"\"\n",
    "    t = df_sh[df_sh['code']==x.iloc[0,1]]['上市日期'].values\n",
    "    if t.size != 0:\n",
    "        condition = x['date'].values > t\n",
    "        x = x[condition]\n",
    "    else:\n",
    "        x = x.loc[x['上市状态_Listedstate'].isin(['Norm', '*ST'])]\n",
    "    return x\n",
    "\n",
    "\n",
    "def Stock_Pool_split(df):\n",
    "    \"\"\"\n",
    "    股票池筛选：\n",
    "    (1) 剔除上市不足1年\n",
    "    (2) 剔除ST\n",
    "    (3) 剔除每股净资产为负\n",
    "    \"\"\"\n",
    "\n",
    "    ## 根据'code'列分组，并将每个分组的DataFrame添加到列表中\n",
    "    grouped_dfs = [group_df for _, group_df in df.groupby('code')]\n",
    "    df_sh = get_sh_time()\n",
    "    y = [] #记录每支股票的信息\n",
    "    l = [] #记录每只股票样本长度\n",
    "    ST = [] # 记录是否被ST过\n",
    "    neg_res = [] # 记录净资产是否为负\n",
    "    hs300 = [] # 记录是否沪深300的股票\n",
    "\n",
    "    for x in grouped_dfs:\n",
    "\n",
    "        ## 去除上市前的样本\n",
    "        x = get_sh_process(x,df_sh)\n",
    "        ## 记录样本长度\n",
    "        l.append(x.shape[0])\n",
    "        ## 记录是否被ST\n",
    "        con_st = x['上市状态_Listedstate'].str.contains('ST').any()\n",
    "        ST.append(con_st)\n",
    "        ## 记录净资产是否为负过\n",
    "        num = (x['每股净资产(元/股)_NAPS']<0).sum()\n",
    "        neg_res.append(num>0)\n",
    "        # ## 记录是否沪深300的股票\n",
    "        if x.shape[0] == 0:\n",
    "            con_hs300 = False\n",
    "        else:\n",
    "            con_hs300 = x['code'].values[0][:6] in hs300_list\n",
    "        hs300.append(con_hs300)\n",
    "\n",
    "\n",
    "        y.append(x)\n",
    "    \n",
    "    ## 剔除上市小于1年的、被ST的、净资产负过的股票\n",
    "    condition = [con_1 and not con_2 and not con_3 and con_4 for con_1, con_2,con_3, con_4 in zip(list(np.array(l) >=365), ST, neg_res,hs300)]\n",
    "    grouped_sel = [item for item, con in zip(y,condition) if con]\n",
    "\n",
    "    return  pd.concat(grouped_sel, ignore_index=True)\n",
    "\n",
    "\n",
    "def Standard_process(df):\n",
    "    \"\"\"\n",
    "    对输入的所有股票的df进行如下操作:\n",
    "    (1) 对每只股票进行：极端值、缺失值处理\n",
    "    (2) 对每只股票的数据进行标准化\n",
    "    \"\"\"\n",
    "    \n",
    "    ## 根据'code'列分组，并将每个分组的DataFrame添加到列表中\n",
    "    grouped_sel= [group_df for _, group_df in df.groupby('code')]\n",
    "\n",
    "    y=[]\n",
    "    for x in grouped_sel:\n",
    "        ## 替换极端值\n",
    "        x.iloc[:,4:x.shape[1]] = extreme_MAD(x.iloc[:,4:x.shape[1]],n = 3)\n",
    "        ## 填补缺失值\n",
    "        x = x.apply(lambda x: x.fillna(method='ffill'))\n",
    "        ## z_score 标准化\n",
    "        scaler = StandardScaler()\n",
    "        x.iloc[:,4:] = scaler.fit_transform(x.iloc[:,4:])\n",
    "\n",
    "        y.append(x)\n",
    "\n",
    "    return pd.concat(y, ignore_index=True)\n",
    "\n",
    "def get_y(df):\n",
    "    \n",
    "    pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('../量化data/股票整合数据_未预处理.zip', compression='zip')\n",
    "df['date'] = pd.to_datetime(df['date'])\n",
    "df_1 = Stock_Pool_split(df)\n",
    "# grouped_dfs = [group_df for _, group_df in df_1.groupby('code')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>date</th>\n",
       "      <th>stock_code</th>\n",
       "      <th>open</th>\n",
       "      <th>high</th>\n",
       "      <th>low</th>\n",
       "      <th>close</th>\n",
       "      <th>factor</th>\n",
       "      <th>change</th>\n",
       "      <th>volume</th>\n",
       "      <th>money</th>\n",
       "      <th>流通股</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2021-02-01</td>\n",
       "      <td>000001.SZ</td>\n",
       "      <td>21.8838</td>\n",
       "      <td>23.7772</td>\n",
       "      <td>21.5984</td>\n",
       "      <td>23.3586</td>\n",
       "      <td>122.9998</td>\n",
       "      <td>0.0632</td>\n",
       "      <td>147523930.0</td>\n",
       "      <td>3.529557e+09</td>\n",
       "      <td>1.940576e+10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2021-02-02</td>\n",
       "      <td>000001.SZ</td>\n",
       "      <td>22.2454</td>\n",
       "      <td>22.9304</td>\n",
       "      <td>21.7982</td>\n",
       "      <td>22.1502</td>\n",
       "      <td>122.9998</td>\n",
       "      <td>-0.0517</td>\n",
       "      <td>241616877.0</td>\n",
       "      <td>5.679180e+09</td>\n",
       "      <td>1.940576e+10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2021-02-03</td>\n",
       "      <td>000001.SZ</td>\n",
       "      <td>22.3405</td>\n",
       "      <td>23.9580</td>\n",
       "      <td>22.2454</td>\n",
       "      <td>23.7392</td>\n",
       "      <td>122.9998</td>\n",
       "      <td>0.0717</td>\n",
       "      <td>192327159.0</td>\n",
       "      <td>4.690176e+09</td>\n",
       "      <td>1.940576e+10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2021-02-04</td>\n",
       "      <td>000001.SZ</td>\n",
       "      <td>23.0066</td>\n",
       "      <td>24.0151</td>\n",
       "      <td>22.8734</td>\n",
       "      <td>23.4062</td>\n",
       "      <td>122.9998</td>\n",
       "      <td>-0.0140</td>\n",
       "      <td>125524750.0</td>\n",
       "      <td>3.084554e+09</td>\n",
       "      <td>1.940576e+10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2021-02-05</td>\n",
       "      <td>000001.SZ</td>\n",
       "      <td>23.4062</td>\n",
       "      <td>24.0817</td>\n",
       "      <td>23.0922</td>\n",
       "      <td>23.7202</td>\n",
       "      <td>122.9998</td>\n",
       "      <td>0.0134</td>\n",
       "      <td>101557559.0</td>\n",
       "      <td>2.517804e+09</td>\n",
       "      <td>1.940576e+10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183242</th>\n",
       "      <td>2023-12-25</td>\n",
       "      <td>605499.SH</td>\n",
       "      <td>176.9300</td>\n",
       "      <td>177.3000</td>\n",
       "      <td>172.6900</td>\n",
       "      <td>175.6200</td>\n",
       "      <td>1.0318</td>\n",
       "      <td>-0.0044</td>\n",
       "      <td>737300.0</td>\n",
       "      <td>1.290807e+08</td>\n",
       "      <td>1.592709e+08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183243</th>\n",
       "      <td>2023-12-26</td>\n",
       "      <td>605499.SH</td>\n",
       "      <td>174.8700</td>\n",
       "      <td>176.9500</td>\n",
       "      <td>174.3400</td>\n",
       "      <td>175.6300</td>\n",
       "      <td>1.0318</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>453300.0</td>\n",
       "      <td>7.969934e+07</td>\n",
       "      <td>1.592709e+08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183244</th>\n",
       "      <td>2023-12-27</td>\n",
       "      <td>605499.SH</td>\n",
       "      <td>175.6800</td>\n",
       "      <td>178.4700</td>\n",
       "      <td>175.6700</td>\n",
       "      <td>177.5500</td>\n",
       "      <td>1.0318</td>\n",
       "      <td>0.0109</td>\n",
       "      <td>671856.0</td>\n",
       "      <td>1.190336e+08</td>\n",
       "      <td>1.592709e+08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183245</th>\n",
       "      <td>2023-12-28</td>\n",
       "      <td>605499.SH</td>\n",
       "      <td>177.9500</td>\n",
       "      <td>178.4700</td>\n",
       "      <td>174.8000</td>\n",
       "      <td>176.9600</td>\n",
       "      <td>1.0318</td>\n",
       "      <td>-0.0033</td>\n",
       "      <td>808121.0</td>\n",
       "      <td>1.425958e+08</td>\n",
       "      <td>1.592709e+08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183246</th>\n",
       "      <td>2023-12-29</td>\n",
       "      <td>605499.SH</td>\n",
       "      <td>176.9600</td>\n",
       "      <td>182.6600</td>\n",
       "      <td>176.6300</td>\n",
       "      <td>182.5100</td>\n",
       "      <td>1.0318</td>\n",
       "      <td>0.0314</td>\n",
       "      <td>800107.0</td>\n",
       "      <td>1.445808e+08</td>\n",
       "      <td>1.592709e+08</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>183247 rows × 11 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             date stock_code      open      high       low     close  \\\n",
       "0      2021-02-01  000001.SZ   21.8838   23.7772   21.5984   23.3586   \n",
       "1      2021-02-02  000001.SZ   22.2454   22.9304   21.7982   22.1502   \n",
       "2      2021-02-03  000001.SZ   22.3405   23.9580   22.2454   23.7392   \n",
       "3      2021-02-04  000001.SZ   23.0066   24.0151   22.8734   23.4062   \n",
       "4      2021-02-05  000001.SZ   23.4062   24.0817   23.0922   23.7202   \n",
       "...           ...        ...       ...       ...       ...       ...   \n",
       "183242 2023-12-25  605499.SH  176.9300  177.3000  172.6900  175.6200   \n",
       "183243 2023-12-26  605499.SH  174.8700  176.9500  174.3400  175.6300   \n",
       "183244 2023-12-27  605499.SH  175.6800  178.4700  175.6700  177.5500   \n",
       "183245 2023-12-28  605499.SH  177.9500  178.4700  174.8000  176.9600   \n",
       "183246 2023-12-29  605499.SH  176.9600  182.6600  176.6300  182.5100   \n",
       "\n",
       "          factor  change       volume         money           流通股  \n",
       "0       122.9998  0.0632  147523930.0  3.529557e+09  1.940576e+10  \n",
       "1       122.9998 -0.0517  241616877.0  5.679180e+09  1.940576e+10  \n",
       "2       122.9998  0.0717  192327159.0  4.690176e+09  1.940576e+10  \n",
       "3       122.9998 -0.0140  125524750.0  3.084554e+09  1.940576e+10  \n",
       "4       122.9998  0.0134  101557559.0  2.517804e+09  1.940576e+10  \n",
       "...          ...     ...          ...           ...           ...  \n",
       "183242    1.0318 -0.0044     737300.0  1.290807e+08  1.592709e+08  \n",
       "183243    1.0318  0.0001     453300.0  7.969934e+07  1.592709e+08  \n",
       "183244    1.0318  0.0109     671856.0  1.190336e+08  1.592709e+08  \n",
       "183245    1.0318 -0.0033     808121.0  1.425958e+08  1.592709e+08  \n",
       "183246    1.0318  0.0314     800107.0  1.445808e+08  1.592709e+08  \n",
       "\n",
       "[183247 rows x 11 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "columns = ['date', 'code','Open', 'High', 'Low', 'Close','累积股价调整乘子_Mcfacpr','日收益率_Dret',\n",
    "           '成交量_Trdvol','成交金额_Trdsum','已上市流通股_Lsttrdshr']\n",
    "df_2 = df_1[columns]\n",
    "df_2 = df_2.rename(columns={'code':'stock_code','Open':'open', 'High':'high', 'Low':'low', 'Close':'close',\n",
    "                            '累积股价调整乘子_Mcfacpr':'factor','日收益率_Dret':'change','成交量_Trdvol':'volume',\n",
    "                            '成交金额_Trdsum':'money','已上市流通股_Lsttrdshr':'流通股'})\n",
    "df_2 \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_2.to_csv('../量化data/df_hs.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## 输出每只股票的csv\n",
    "# codes = df_2['stock_code'].unique()\n",
    "# for i in codes:\n",
    "#     df = df_2.loc[df_2['stock_code'] == i]\n",
    "    \n",
    "#     #将每支股票的相关信息分别导出成一个csv文件\n",
    "#     df.to_csv('C:/Users/Lincoln/.qlib/csv_data/ch_data/'+ i +'.csv',header = True,encoding='utf_8_sig',index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python scripts/dump_bin.py dump_all --csv_path ~/.qlib/csv_data/ch_data --qlib_dir ~/.qlib/qlib_data/ch_data --symbol_field_name stock_code --date_field_name date --include_fields name,volume,money,factor,close,open,high,low,change"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13 (main, Aug 25 2022, 23:51:50) [MSC v.1916 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f52f256adf2cc158a51e888fcb9b6688be269c3160efd3d6e246b7ec509c991c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
